#!/usr/bin/env python3
"""
数据库迁移脚本：将class_name字段改为class_id整型字段
修改字段类型从VARCHAR改为INTEGER，并添加约束(0-10)
"""

import sqlite3
import os
import sys

def get_db_path():
    """获取数据库文件路径"""
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    return os.path.join(base_dir, 'instance', 'mailserver_dev.db')

def upgrade():
    """
    升级数据库：将class_name字段改为class_id整型字段
    """
    db_path = get_db_path()
    
    if not os.path.exists(db_path):
        print(f"错误：数据库文件不存在: {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        print("开始数据库迁移：将class_name改为class_id...")
        
        # 1. 检查是否存在class_name字段
        cursor.execute("PRAGMA table_info(users)")
        columns = cursor.fetchall()
        column_names = [col[1] for col in columns]
        
        if 'class_name' not in column_names:
            print("class_name字段不存在，跳过迁移")
            conn.close()
            return True
        
        # 2. 创建新表结构
        cursor.execute("""
            CREATE TABLE users_new (
                id INTEGER PRIMARY KEY,
                username VARCHAR(80) NOT NULL UNIQUE,
                email VARCHAR(120) NOT NULL UNIQUE,
                password_hash VARCHAR(255) NOT NULL,
                role VARCHAR(20) NOT NULL DEFAULT 'student',
                real_name VARCHAR(50),
                class_id INTEGER CHECK (class_id >= 0 AND class_id <= 10),
                is_active BOOLEAN NOT NULL DEFAULT 1,
                created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                last_login DATETIME
            )
        """)
        
        # 3. 创建索引
        cursor.execute("CREATE INDEX ix_users_new_username ON users_new (username)")
        cursor.execute("CREATE INDEX ix_users_new_email ON users_new (email)")
        cursor.execute("CREATE INDEX ix_users_new_class_id ON users_new (class_id)")
        
        # 4. 迁移数据，将class_name转换为class_id
        # 这里需要根据实际的class_name值来映射到class_id
        # 假设class_name格式为"计算机科学1班"等，提取数字作为class_id
        cursor.execute("""
            INSERT INTO users_new (
                id, username, email, password_hash, role, real_name, 
                class_id, is_active, created_at, last_login
            )
            SELECT 
                id, username, email, password_hash, role, real_name,
                CASE 
                    WHEN class_name IS NULL OR class_name = '' THEN NULL
                    WHEN class_name LIKE '%1%' THEN 1
                    WHEN class_name LIKE '%2%' THEN 2
                    WHEN class_name LIKE '%3%' THEN 3
                    WHEN class_name LIKE '%4%' THEN 4
                    WHEN class_name LIKE '%5%' THEN 5
                    WHEN class_name LIKE '%6%' THEN 6
                    WHEN class_name LIKE '%7%' THEN 7
                    WHEN class_name LIKE '%8%' THEN 8
                    WHEN class_name LIKE '%9%' THEN 9
                    WHEN class_name LIKE '%10%' THEN 10
                    ELSE 0
                END as class_id,
                is_active, created_at, last_login
            FROM users
        """)
        
        # 5. 删除旧表
        cursor.execute("DROP TABLE users")
        
        # 6. 重命名新表
        cursor.execute("ALTER TABLE users_new RENAME TO users")
        
        conn.commit()
        print("✅ 数据库迁移完成：class_name已改为class_id")
        print("✅ 添加了约束：class_id范围为0-10")
        
        # 显示迁移结果
        cursor.execute("SELECT COUNT(*) FROM users WHERE class_id IS NOT NULL")
        count = cursor.fetchone()[0]
        print(f"✅ 成功迁移了 {count} 个用户的班级信息")
        
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ 数据库迁移失败: {str(e)}")
        if 'conn' in locals():
            conn.rollback()
            conn.close()
        return False

def downgrade():
    """
    降级数据库：将class_id字段改回class_name字符串字段
    """
    db_path = get_db_path()
    
    if not os.path.exists(db_path):
        print(f"错误：数据库文件不存在: {db_path}")
        return False
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        print("开始数据库降级：将class_id改回class_name...")
        
        # 1. 检查是否存在class_id字段
        cursor.execute("PRAGMA table_info(users)")
        columns = cursor.fetchall()
        column_names = [col[1] for col in columns]
        
        if 'class_id' not in column_names:
            print("class_id字段不存在，跳过降级")
            conn.close()
            return True
        
        # 2. 创建新表结构
        cursor.execute("""
            CREATE TABLE users_new (
                id INTEGER PRIMARY KEY,
                username VARCHAR(80) NOT NULL UNIQUE,
                email VARCHAR(120) NOT NULL UNIQUE,
                password_hash VARCHAR(255) NOT NULL,
                role VARCHAR(20) NOT NULL DEFAULT 'student',
                real_name VARCHAR(50),
                class_name VARCHAR(50),
                is_active BOOLEAN NOT NULL DEFAULT 1,
                created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
                last_login DATETIME
            )
        """)
        
        # 3. 创建索引
        cursor.execute("CREATE INDEX ix_users_new_username ON users_new (username)")
        cursor.execute("CREATE INDEX ix_users_new_email ON users_new (email)")
        cursor.execute("CREATE INDEX ix_users_new_class_name ON users_new (class_name)")
        
        # 4. 迁移数据，将class_id转换回class_name
        cursor.execute("""
            INSERT INTO users_new (
                id, username, email, password_hash, role, real_name, 
                class_name, is_active, created_at, last_login
            )
            SELECT 
                id, username, email, password_hash, role, real_name,
                CASE 
                    WHEN class_id IS NULL THEN NULL
                    WHEN class_id = 0 THEN '其他班级'
                    WHEN class_id = 1 THEN '计算机科学1班'
                    WHEN class_id = 2 THEN '计算机科学2班'
                    WHEN class_id = 3 THEN '软件工程1班'
                    WHEN class_id = 4 THEN '软件工程2班'
                    WHEN class_id = 5 THEN '数据科学1班'
                    WHEN class_id = 6 THEN '网络工程1班'
                    WHEN class_id = 7 THEN '信息安全1班'
                    WHEN class_id = 8 THEN '人工智能1班'
                    WHEN class_id = 9 THEN '物联网1班'
                    WHEN class_id = 10 THEN '大数据1班'
                    ELSE '未知班级'
                END as class_name,
                is_active, created_at, last_login
            FROM users
        """)
        
        # 5. 删除旧表
        cursor.execute("DROP TABLE users")
        
        # 6. 重命名新表
        cursor.execute("ALTER TABLE users_new RENAME TO users")
        
        conn.commit()
        print("✅ 数据库降级完成：class_id已改回class_name")
        
        conn.close()
        return True
        
    except Exception as e:
        print(f"❌ 数据库降级失败: {str(e)}")
        if 'conn' in locals():
            conn.rollback()
            conn.close()
        return False

if __name__ == '__main__':
    if len(sys.argv) > 1:
        if sys.argv[1] == 'upgrade':
            success = upgrade()
        elif sys.argv[1] == 'downgrade':
            success = downgrade()
        else:
            print("用法: python modify_class_field.py [upgrade|downgrade]")
            sys.exit(1)
    else:
        print("用法: python modify_class_field.py [upgrade|downgrade]")
        sys.exit(1)
    
    sys.exit(0 if success else 1)